'''
$LastChangedBy: c.gillespie@ncl.ac.uk $
$Revision: 123 $
$LastChangedDate: 2009-04-24 17:27:10 +0000 (Fri, 24 Apr 2009) $

If want change the shape/colour/etc of a node or arrow
look at VisualiseNodeClasses
'''
import pydot    
from VisualiseNodeClasses import DotNodes

class VisualiseModel(DotNodes):
    """The visualisation model class"""
    
    def    __init__(self, m, style_dictionary):
        '''
        m is a libsbml getModel file
        initialise a few short-cuts here
        '''
        self.m = m
        self.styles = style_dictionary
        
        graph_attributes = self.styles['graph_attributes']
        self.py_dot_graph = pydot.Dot(**graph_attributes)
        
       
        self.add_edge = self.py_dot_graph.add_edge
        self.add_node = self.py_dot_graph.add_node
        DotNodes.__init__(self, m, self.styles)

    def __createCompartmentDict(self):
        """creates a dictionary with all compartments as keys and a 
            pydot cluster as a value
        """
        compart_dict = {}
        for i, co in enumerate(self.m.getListOfCompartments()):
            compart_dict[co.getId()] = [co.getOutside()]
            node_attribute = self.compartmentNode(i+1)
            compart_dict[co.getId()].append(pydot.Cluster(**node_attribute))
        return compart_dict
    
    def __organiseCompartments(self, compart_dict):
        """organises the compartments, using outside"""
        for key in compart_dict.keys():
            outside = compart_dict[key][0]
            if len(outside) ==0:
                self.py_dot_graph.add_subgraph(compart_dict[key][1])
            else:
                compart_dict[outside][1].add_subgraph(compart_dict[key][1])
        return self.py_dot_graph    

    def __isGraph2Big(self):
        '''Is the model too big 
            XXX: not affected by the user dictionary
        '''
        if (self.m.getNumReactions() + self.m.getNumSpecies()) < 250:
            return 0

        dict1 = {'URL':'#', 'label':'There are too many reactions/ \
                    species to be display'}
        node_attribute = self.compartmentNode(1, dict1)
        cluster = pydot.Cluster(**node_attribute)
        
        node_attribute = self.reactionNode(2, {'URL':'#'})
        node_attribute['label'] = '%d Reactions' % self.m.getNumReactions()
        cluster.add_node(pydot.Node('reactions', **node_attribute))
        
        node_attribute = self.specieNode(2, {'URL':'#'})
        node_attribute['label'] = '%d species' % self.m.getNumSpecies()
        cluster.add_node(pydot.Node('species', **node_attribute))
        
        self.py_dot_graph.add_subgraph(cluster)
        return 1
            
    def __isCompartments(self):
        '''Does the model have any compartments
        XXX: not affected by the user dictionary
        '''
        if self.m.getNumCompartments() > 0:
            return 0
        
        node_attribute = {'graph_name':'empty model', 'label':'empty model', \
                'fontsize':"10", 'fontname':'Arial'}
        node_attribute['URL'] = '#'
        cluster = pydot.Cluster(**node_attribute)
            
        node_attribute = {'color':'moccasin', 'shape':'rectangle',
            'style':'filled', 'label':'0 species', 'fontsize':"10"}
        node_attribute['fontname'] = 'Arial'
        node_attribute['URL'] = '#'
        cluster.add_node(pydot.Node('species', **node_attribute))
    
        self.py_dot_graph.add_subgraph(cluster)
        return 1
    
    def __isSpecies(self):
        '''Does the model have any species
        XXX: not affected by the user dictionary'''
        if self.m.getNumSpecies() > 0:
            return 0
        
        node_attribute = self.compartmentNode(1)
        cluster = pydot.Cluster(**node_attribute)
        node_attribute = {'color':'moccasin', 'shape':'rectangle',
            'style':'filled', 'label':'0 species', 'fontsize':"10"}
        node_attribute['fontname'] = 'Arial'
        node_attribute['URL'] = '#'
        cluster.add_node(pydot.Node('species', **node_attribute))
        
        self.py_dot_graph.add_subgraph(cluster)
        return 1
    
    def __addReactionToCompartment(self, sp_ref):
        '''Have to put the reaction into a compartment so
        choose the first reactant and stick it that compartment
        If there are no reactants than get the first product
        '''
        for sp in self.m.getListOfSpecies():
            if sp.getId() == sp_ref.getSpecies():
                return 1
    
    def getPyDot(self):
        '''main method'''
        #If there are too many reactions
        
        if self.__isGraph2Big():
            return self.py_dot_graph
        
        if self.__isCompartments():
            return self.py_dot_graph
            
        if self.__isSpecies():
            return self.py_dot_graph

        compart_dict = self.__createCompartmentDict()
        #Add the species into a compartment
        for i, sp in enumerate(self.m.getListOfSpecies()):
            node_attribute = self.specieNode(i+1)
            
            compart_dict[sp.getCompartment()][1].add_node(
                pydot.Node(sp.getId(), **node_attribute))
    
        for i, reaction in enumerate(self.m.getListOfReactions()):
            node_attribute = self.reactionNode(i+1)
            for j, sp_ref in enumerate(reaction.getListOfReactants()):
                if j == 0 and self.__addReactionToCompartment(sp_ref):
                    compart_dict[sp.getCompartment()][1].add_node(
                        pydot.Node(reaction.getId(), **node_attribute))
                
                stoic = sp_ref.getStoichiometry()
                #print type(stoic)
                if divmod(stoic, 1)[1] > 0.000001:#Is the stoic an integer
                    stoic = str(stoic)
                        
                edge_attribute = self.edgeNode(reaction, stoic, 'reactant_edge')
                edge = pydot.Edge(sp_ref.getSpecies(), reaction.getId(), **edge_attribute)
                self.add_edge(edge)
                
            for j, sp_ref in enumerate(reaction.getListOfProducts()):
                if j == 0 and not reaction.getNumReactants() and \
                    self.__addReactionToCompartment(sp_ref):
                    compart_dict[sp.getCompartment()][1].add_node(
                            pydot.Node(reaction.getId(), **node_attribute))
                        
                stoic = sp_ref.getStoichiometry()
                if divmod(stoic,1)[1] > 0.000001:#Is the stoic an integer
                    stoic = str(stoic)
            
                edge_attribute = self.edgeNode(reaction, stoic, 'product_edge')
                edge = pydot.Edge(reaction.getId(), 
                        sp_ref.getSpecies(), **edge_attribute)
                self.add_edge(edge)
    
            for sp_ref in reaction.getListOfModifiers():
                edge_attribute = self.edgeNode(reaction, 1, 'modifier_edge')
                edge = pydot.Edge(sp_ref.getSpecies(), reaction.getId(), **edge_attribute)
                self.add_edge(edge)
                edge = pydot.Edge(reaction.getId(), sp_ref.getSpecies(), **edge_attribute)
                self.add_edge(edge)

        return self.__organiseCompartments(compart_dict)
    

class VisualiseReaction(DotNodes):
    '''Used to visualise a single reaction'''
    
    def    __init__(self, m, re_num, style_dictionary):
        '''
        m is libsbml.getModel 
        re_num is the reaction to be visualised
        '''
        self.m = m
        self.styles = style_dictionary
        
        graph_attributes = self.styles['graph_attributes']
        self.py_dot_graph = pydot.Dot(**graph_attributes)
       
       
        self.add_edge = self.py_dot_graph.add_edge
        self.add_node = self.py_dot_graph.add_node
             
        self.re_num = re_num
        self.re = self.m.getReaction(re_num-1)
        
        DotNodes.__init__(self, m, self.styles)

    def getPyDot(self):
        '''Returns a pyDot object'''
                
        #Create the reaction node
        node_attribute = self.reactionNode(self.re_num)
        self.add_node(pydot.Node(self.re.getId(), **node_attribute))
            
        types = ('Reactants', 'Products', 'Modifiers')
        for typ in types:
            for sp_ref in getattr(self.re, 'getListOf' + typ)():
                sp_ref_id = sp_ref.getSpecies()
                self.__findMatchingSpecies(sp_ref_id)
            
                if typ in ('Reactants','Products'):
                    stoic = sp_ref.getStoichiometry()
                    if divmod(stoic, 1)[1] > 0.000001:#Is the stoic an integer
                        stoic = str(stoic)
                else:
                    stoic = str(1)
                
                edge_type = typ[0:-1].lower() + '_edge'
                edge_attribute = self.edgeNode(self.re, stoic, edge_type)
                if typ != 'Products':
                    edge = pydot.Edge(sp_ref_id, self.re.getId(), **edge_attribute)
                    self.add_edge(edge)
                if typ != 'Reactants':
                    edge = pydot.Edge(self.re.getId(), sp_ref_id, **edge_attribute)
                    self.add_edge(edge)
        return self.py_dot_graph            
                
    def __findMatchingSpecies(self, sp_ref_id):
        '''determines where the sp_ref is in listOfSpecies'''
        
        #TODO: make a species list, so we don't keep cycling through sp!
        for i, sp in enumerate(self.m.getListOfSpecies()):
            if sp.getId() == sp_ref_id:
                node_attribute = self.specieNode(i+1)
                self.add_node(pydot.Node(sp_ref_id, **node_attribute))
                break
        return True
    

class VisualiseSpecies(DotNodes):
    '''Visualise a single species'''
    
    def    __init__(self, m, sp_num, style_dictionary):
        '''takes in a sbml and validates it'''
        self.m = m
        self.styles = style_dictionary
        
        graph_attributes = self.styles['graph_attributes']
        self.py_dot_graph = pydot.Dot(**graph_attributes)
        self.add_edge = self.py_dot_graph.add_edge
        self.add_node = self.py_dot_graph.add_node
        
        DotNodes.__init__(self, m, self.styles)
        
        self.sp_num = sp_num
        self.sp_id = self.m.getSpecies(sp_num-1).getId()
    
    def getPyDot(self):
        '''If you want to visualise a species'''
        #Create the species node
        node_attribute = self.specieNode(self.sp_num)
        self.add_node(pydot.Node(self.sp_id, **node_attribute))
                
        for i, reaction in enumerate(self.m.getListOfReactions()):
            self.__checkReactionTypesForSpecies(reaction, i)
        return self.py_dot_graph
            
    def __checkReactionTypesForSpecies(self, reaction, reaction_num):
        '''Checks all the species references to determine if the sp is 
        in this reaction
        '''
        types = ('Reactants', 'Products', 'Modifiers')
        for typ in types:
            for sp_ref in getattr(reaction, 'getListOf' + typ)():
                if sp_ref.getSpecies() == self.sp_id:
                    node_attribute = self.reactionNode(reaction_num+1)
                    r_id = reaction.getId()
                    self.add_node(pydot.Node(r_id, **node_attribute))
                    
                    if typ in ('Reactants','Products'):
                        stoic = sp_ref.getStoichiometry()
                        if divmod(stoic, 1)[1] > 0.000001:#Is the stoic an integer
                            stoic = str(stoic)

                    else:
                        stoic = str(1)
                    edge_type = typ[0:-1].lower() + '_edge'
                    edge_attribute = self.edgeNode(reaction, stoic, edge_type)

                    if typ != 'Reactants':
                        self.add_edge(pydot.Edge(r_id, self.sp_id, **edge_attribute))
                    if typ != 'Products':
                        self.add_edge(pydot.Edge(self.sp_id, r_id, **edge_attribute))
                    return 1




